package com.aismy.admin.websocket;

import com.aismy.admin.entity.ConsultationMessage;
import com.aismy.admin.service.ConsultationMessageService;
import com.aismy.admin.utils.Json;
import com.aismy.admin.utils.SecurityUtils;
import com.aismy.admin.utils.SpringContextUtil;
import jakarta.annotation.Resource;
import jakarta.websocket.*;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
@ServerEndpoint("/ws/consultation/{consultationId}/{userId}")
public class ConsultationWebSocketService {

    protected static final Map<String, Map<String, Session>> CONSULTATION_SESSIONS = new ConcurrentHashMap<>();


    @Data
    protected static class ConsultationSession {
        Session userSession;
        Session consultantSession;
    }

    @OnOpen
    public void onOpen(Session session,
                       @PathParam("consultationId") String consultationId,
                       @PathParam("userId") String userId) {
        log.info("新的WebSocket连接建立 consultationId: {}, role: {}", consultationId, userId);

        // 存储或更新咨询会话
        CONSULTATION_SESSIONS.compute(consultationId, (k, v) -> {
            if (v == null) {
                v = new HashMap<>(2);
            }
            v.put(userId, session);
            return v;
        });
    }


    public void sendMessage(ConsultationMessage message) {
        Map<String, Session> consultationSession = CONSULTATION_SESSIONS.get(String.valueOf(message.getConsultationId()));
        if (consultationSession != null) {
            Session session = consultationSession.get(String.valueOf(message.getReceiverId()));
            if (session != null) {
                try {
                    session.getAsyncRemote().sendText(Json.toJson(message));
                } catch (Exception e) {
                    log.error("发送消息失败", e);
                }
            }
        }
    }

    @OnClose
    public void onClose(Session session,
                        @PathParam("consultationId") String consultationId,
                        @PathParam("userId") String userId) {
        log.info("WebSocket连接关闭 consultationId: {}, role: {}", consultationId, userId);

        // 更新咨询会话
        Map<String, Session> consultationSession = CONSULTATION_SESSIONS.get(consultationId);
        if (consultationSession != null) {
            Session remove = consultationSession.remove(userId);
            if (remove != null) {
                try {
                    remove.close();
                } catch (IOException e) {
                    log.error("关闭WebSocket连接失败", e);
                }
            }
            // 如果双方都断开连接，则移除咨询会话
            if (consultationSession.isEmpty()) {
                CONSULTATION_SESSIONS.remove(consultationId);
            }
        }
    }
} 